Refactored ks_twosample to take in two slices instead of a mutable owned Vec<f64> - #406

@AshrafIbrahim03 did you verify tests with all targets n features?

I did not. I just reread through the contributing section of the README to see if it details how to run those, but I don't see it. How can I do that?

probably you find out:
cargo test

Thanks, the call sites compile now. I checked the latest commit but three KS tests still fail. currently order is wrong. are you sure it should be like that? may you add some context to Pr ?

Yes, just saw the last three tests fail, pushing a fix for it now!

Some more context: basically we were talking about a contribution I could make in #405 , and it seemed like changing the ks sampling function to take a &[f64] instead of a mut Vec<f64> would be more in line with the other files, anderson_darling and chi_square.

I also noticed the ks one sample function that I should probably bundle in this PR too!

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.82%. Comparing base (ad9676b) to head (0f070d0).

@@           Coverage Diff           @@
##             main     #406   +/-   ##
=======================================
  Coverage   94.82%   94.82%           
=======================================
  Files          61       61           
  Lines       13539    13539           
=======================================
  Hits        12838    12838           
  Misses        701      701           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

I realize I steered you astray. I think it's better to let the caller opt into copying data. Ideally, we could obtain an iterator over the "sorted" data borrowing, but I don't know if that's easy to express, or we just verify order along with nan policy with an iterator argument.

Thoughts as a user?

If I need to retain the input data, I can clone it explicitly, otherwise the function can consume and sort it without hidden copies

@AshrafIbrahim03 that means we should go the other way, convert those that accept slices to instead own as Vec for data we sort in place.

It would be great if there were a way to get an iterator that's element-wise sorted, because then the API expresses we need each value in order. It's incidental that we happen to sort the data ourselves to access that ordering of the values efficiently. Maybe we'll defer that for larger datasets that need to be streamed.

@AshrafIbrahim03 need any help here?

@AshrafIbrahim03 need any help here?

Just coming back to this now.

It would be great if there were a way to get an iterator that's element-wise sorted, because then the API expresses we need each value in order.

This sounds like it could be a new interface that implements Iterator, but calling next just returns the sorted elements, no?

Accidentally overwrote my prior commits, but it didn't seem like those were needed. I pushed a commit that has some code with a basic sorted iterator. The implementation is not efficient, but is that the type of input you're looking for in the ks_test functions?

i dont think so, it still will be clone.

The implementation is not efficient, but is that the type of input you're looking for in the ks_test functions?

Think I'd need to see the API for the ks_*sample functions to be certain, but I do agree with @day01 that if you wanted to provide the Sorted as is, it would require a clone since it's a data structure type (owns structure and values) since there's not a Rust standard notion of sorted iterator, it also might require a constructor that doesn't sort in case I don't want my data explicitly sorted again (perhaps it was sorted on write to disk, and read protocol will read sorted) and we just validate order on the fly.

But we have some good constraints, we need to be able to have a pull-iterator in a sorted order (the for loop in the algorithm) and we want to avoid a clone of the data in the source structure.

Let me know if you're okay with hints/want something clearer/have argument we're missing for your approach. If you write out a usage example that exhibits both of these:

• calls the ks_*sample function that has API expressing, "I have provided an iterator that I assert is sorted"
• while avoiding a struct/type that the user has to handle explicitly (sim to most structs in core::iter like Chain, Cloned, we don't have to think about them)

It could also be a good simplifying point to assume that you start with a Vec<T: Ord> and then generalize from there asking "what would I need to provide to express a similar constraint?"

I think you're right that taking in a new data structure, like Sorted shouldn't be the move. I'm wondering if taking in an impl IntoSortedIterator would be a better way to go about this, then implementing that type for different types of collections would allow a caller more flexibility with the passed argument. I'm thinking the interface of interacting with sorted data could be similar to parallelizing using rayon's parallelized iterator? Simply calling IntoSortedIterator::into_sorted_iterator on an Iterator<Item=T:Ord> would return a SortedIterator, allowing the ks_*sample to interact with these collections without mutating the vector itself.

I think an Iterator implementation would be the best approach, because only reason for mutability in the functions is for sorting, even calculating the test statistics takes iterators

There's one problem with this approach that I've been looking into during my spare time, that being how to iterate through a collection in a sorted order without mutating the underlying data structure or making the iterator itself expensive. I've been doing some research trying to figure this out with just Vec<f64> before generalizing, but if you have any guidance here that would be great!

I'm working on a minimum viable rewrite of the ks_*sample functions using data: impl IntoSortedIterator just to see if this change would be feasible in the sampling methods, but after evaluating feasibility there I would figure out how to optimize the Sorted to not just clone and sort the input Vector.

Let me know if this is a good approach or not. If this isn't I would like some more guidance on how to approach this problem!

• Fix all pre-merge checks with AI

_{Comment @coderabbitai help to get the list of available commands.}

I got stuck on this too. The difficult part seems to be that the property we want to express is about the values yielded by the iterator, rather than something its type necessarily enforces. A structure like a B-tree can guarantee ordering as an invariant, but an arbitrary iterator generally can't.

There's also a stronger assumption hiding here: the iterator is traversing a fixed dataset. That's more restrictive than just ExactSizeIterator, but it still doesn't seem sufficient to statically establish that the yielded values are ordered. Unless the data is already known to be sorted, I think implementing such an API directly would be difficult, so how can we express or mark that the data has already been sorted?

I thought about what you said a little bit and realized that you meant making sorting the caller's issue. I made a new struct called SortedCollection in sorted_collection.rs that verifies that an input slice is sorted on instantiation and fails if it isn't sorted. The struct itself exposes the iter method on its inner slice so that the inner slice has been validated and an iterator over it can be called when needed.

When I was looking at your hints, I thought that they leaned towards Haskell style arrays with compile time bounds on the values. The issue I see with that approach would be that it would mean that the sorted verification could only happen at compile time. If that's not what you meant by it, I would love more clarification so I can try and implement it.

I'm also thinking that the ks_*sample methods could take a borrow to the SortedCollection struct so that any calculations are done with an iterator derived from a verified sorted collection. If this is the implementation we want to go towards, I think that have some helper methods with making a SortedCollection would go a long way to help with usability.

Yes, I thought to put it on the caller. The compile level I mentioned would only have the value of a marker trait so would since we can't figure out if data is sorted by its type unless it's from a data structure that has a traversal that guarantees some ordering.

I think SortedCollection is a good start, perhaps named SortedSlice for specificity, and it doesn't seem close to widen the ks two sample API into something like an iterator extension or newtype that SortedSlice provides.

The reason for an iterator API is that I know someone is using statrs for a polars plugin, and polars handles upstream sorts in an execution plan by a marker variable, so there's not actually data yet let alone sorted, but by the time statrs algorithms are executing there will be sorted data and by then I think it will also store the size as well. @FBruzzesi do you have any input on how you'd consume this API for your polars plugin if you were to use it?

Thanks for the ping @YeungOnion

Context for my use case first. Today polars-stats only uses statrs::distribution, and hypothesis tests are out of scope for now. That said, it might be nice to eventually include them.

Disclaimer: I have not read the polars plugin internals, so treat the below as observed behaviour rather than the documented contract. The plugin API is thinly documented, so it is worth double-checking with the polars team (they have a Discord with dedicated plugins and rust channels).

My understanding is that by the time my plugin function runs, polars has already planned the query and executed the sort. What arrives is plain materialized data, and there is no laziness left on the Rust side of the plugin boundary. I have also never found a way to ask, from inside a plugin, whether the column was sorted upstream.

Even if I could check it, polars' notion of sorted is a looser promise than SortedCollection makes. A column flagged sorted ascending can have nulls at the front and NaN at the back, and a user can set that flag by hand with no validation at all:

>>> pl.Series([3.0, float("nan"), None, 1.0]).sort().to_list()
[None, 1.0, 3.0, nan]                      # flagged SORTED_ASC

>>> pl.Series([3.0, 1.0, 2.0]).set_sorted().flags
{'SORTED_ASC': True, 'SORTED_DESC': False} # not sorted, no complaint

So I would validate on the plugin side regardless, and that is fine (it's one linear scan). The two things that do cost are elsewhere.

1. Lend the data instead of handing it over: The values live in memory polars owns and the plugin borrows. A Vec<f64> parameter means copying the whole column before the test even starts. The &[f64] in this PR is the biggest win, independent of the sorted question.
2. Do not require the column to be one contiguous block: Polars can hold a column in several chunks, and I only get a flat slice when there is exactly one chunk and no nulls. So a slice-only API makes me glue the chunks back together first, which is the full copy the slice was meant to avoid. This is the real argument for an iterator API: not deferred data, just data in pieces.

Concretely, the difference would look like something along these lines:

let a = a_series.f64()?.drop_nulls();  // polars nulls are not NaN; NaNPolicy still applies to NaN
let b = b_series.f64()?.drop_nulls();

// slice API: a full copy each whenever the column arrived in more than one chunk
let (a, b) = (a.rechunk(), b.rechunk());
let out = ks_twosample(
    SortedSlice::try_from_slice(a.cont_slice()?)?,
    SortedSlice::try_from_slice(b.cont_slice()?)?,
    method,
    nan_policy,
)?;

// iterator API: no copy, chunks walked in place
let out = ks_twosample(
    a.into_no_null_iter().assert_sorted(),
    b.into_no_null_iter().assert_sorted(),
    method,
    nan_policy,
)?;

So to answer your question directly: SortedSlice is the one I would reach for first and it is already a clear improvement over the owned Vec, but the iterator version is the one that would actually save a copy.

I hope this helps 😇